from ._window import Window
from ._window_data import PartialWindowData
import torch


class CountWindow(Window):
    def __init__(self, data_loader, window_size, threshold):
        super().__init__()
        self.data_loader = iter(data_loader)
        self.window_size = window_size
        self.threshold = threshold
        self.offset = 0

    def __inner_poll(self):
        try:
            X, y = next(self.data_loader)
        except StopIteration:
            return None
        self.offset += len(X)
        trigger_window = False
        if self.offset >= self.window_size:
            self.offset = 0
            trigger_window = True
        window_data = PartialWindowData(X, y, trigger_window)
        return window_data

    def poll(self) -> PartialWindowData:
        return self.__inner_poll()
